/********** System Includes ***************************************************/

#include <iostream>
#include <functional>
#include <list>

/********** Local Includes ****************************************************/

#include "blobdetector.h"

#include "common.h"

/********** Macros and Constants **********************************************/

#define RED 2
#define GREEN 1
#define BLUE 0

/********** Namespaces ********************************************************/

using namespace cv;
using namespace std;

/********** Local Static Function Prototypes **********************************/

bool blobCompare(const BlobDetector::Blob &one, const BlobDetector::Blob &two);
double euclidean(Point p1, Point p2);

/********** Public Methods ****************************************************/

BlobDetector::BlobDetector() {
	SimpleBlobDetector::Params params;
	params.minThreshold = 40;
	params.maxThreshold = 60;
	params.thresholdStep = 5;
	params.minArea = 50; 
	params.minConvexity = 0;
	params.minInertiaRatio = 0;
	params.maxArea = 60000;
	params.maxConvexity = 255;
	params.filterByColor = false;
	params.filterByCircularity = false;
	
	m_sbd = SimpleBlobDetector(params);
	m_sbd.create("SimpleBlob");
	
	m_colors.push_back(COLOR_BLUE);
	m_colors.push_back(COLOR_GREEN);
}

vector<BlobDetector::Blob> BlobDetector::findBlobs(const Mat &image) {
	vector<Blob> blobs;
	
	if (image.empty()) {
		ERROR("no image to work with");
		return vector<Blob>();
	}
	
	// Iterates over all of the elements in m_colors
	for (ColorIterator it = m_colors.begin(); it != m_colors.end(); ++it) {
		std::vector<Blob> result = findBlobs(image, *it);
		
		// Adds all of the elements in result to blobs
		blobs.insert(blobs.end(), result.begin(), result.end());
	}
	
	return blobs;
}

vector<BlobDetector::Blob> BlobDetector::filterForLargest(
	const vector<BlobDetector::Blob> &allBlobs) const
{
	map<Color, int> max; // The maximum number of blobs to accept for each color
	map<Color, int> accepted; // The number already accepted
	
	for (ColorIterator it = m_colors.begin(); it != m_colors.end(); ++it)
	{
		// By default, only accept one of each color
		max[*it] = 1;
		accepted[*it] = 0;
	}
	
	// We'll sort the blobs by size, then take as many of each color as we can
	vector<Blob> copy(allBlobs);
	sort(copy.begin(), copy.end(), blobCompare);
	
	vector<Blob> largest;
	
	for (vector<Blob>::const_iterator it = copy.begin(); it != copy.end(); ++it) {
		map<Color, int>::const_iterator iMax = max.find(it->color);
		map<Color, int>::const_iterator iAccepted = accepted.find(it->color);
		
		if (iMax == max.end() || iAccepted == accepted.end()) {
			WARNING("allBlobs contains a blob with an unrecognized color");
			continue;
		}
		
		// If we have not yet reached the quota for this color...
		if (iMax->second > iAccepted->second) {
			largest.push_back(*it);
			accepted[it->color]++;
		}
	}
	
	return largest;
}

map<BlobDetector::BlobCategory, BlobDetector::Blob> BlobDetector::categorize(
	const vector<BlobDetector::Blob> &largest) const
{
	map<BlobCategory, Blob> cat;
	
	for (vector<Blob>::const_iterator it = largest.begin();
		it != largest.end(); ++it)
	{
		if (it->color == COLOR_GREEN && cat.find(ROBOT_RIGHT) == cat.end()) {
			cat[ROBOT_RIGHT] = *it;
			cat[ROBOT_RIGHT].category = ROBOT_RIGHT;
		} else if (it->color == COLOR_BLUE && cat.find(ROBOT_LEFT) == cat.end()) {
			cat[ROBOT_LEFT] = *it;
			cat[ROBOT_LEFT].category = ROBOT_LEFT;
		}
	}
	
	return cat;
}

map<BlobDetector::BlobCategory, Point> BlobDetector::getPoints(const Mat &image) {
	vector<Blob> blobs = findBlobs(image);
	blobs = filterForLargest(blobs);
	
	map<BlobCategory, BlobDetector::Blob> cat = categorize(blobs);
	map<BlobCategory, Point> points;
	
	for (map<BlobCategory, Blob>::const_iterator it = cat.begin();
		it != cat.end(); ++it)
	{
		points[it->first] = it->second.center;
	}
	
	return points;
}

bool any(const Vec3b &vec, int threshold = 0) {
	return vec[0] > threshold || vec[1] > threshold || vec[2] > threshold;
}

vector<Point> neighbors(Point p, Size s) {
	static int d[4][2] = {
		{ 1, 0 },
		{ -1, 0 },
		{ 0, 1 },
		{ 0, -1 }
	};
	
	vector<Point> neighbors;
	
	for (int i = 0; i < 4; i++) {
		Point c(p.x + d[i][0], p.y + d[i][1]);
		
		if (c.x >= 0 && c.y >= 0 && c.x < s.width && c.y < s.height) {
			neighbors.push_back(c);
		}
	}
	
	return neighbors;
}

int hash(const Point &p) {
	return p.x * 1000 + p.y;
}

vector<Point> BlobDetector::pixels(const Mat &image,
	const BlobDetector::Blob &blob) const
{
	Mat filtered;
	filter(image, filtered, blob.color);
	
	vector<Point> pixels;
	list<Point> candidates;  // Used as a queue because std::queue sucks
	
	/*
	 * Contains a hash of each point to quickly determine if a point has already
	 * been added.
	 *
	 * Either the STL doesn't have an ordered list implementation, or I missed
	 * it in the documentation.  A data structure with O(log n) insertion and
	 * look-up is criticial; the performance advantage over a structure with
	 * O(n) look-up is on the order of seconds.
	 */
	map<int, int> hashes;
	
	candidates.push_back(blob.center);
	
	while (!candidates.empty()) {
		Point candidate = candidates.front();
		candidates.pop_front();
		
		if (any(filtered.at<Vec3b>(candidate.y, candidate.x), 10)
			&& !hashes.count(hash(candidate)))
		{
			pixels.push_back(candidate);
			hashes[hash(candidate)] = 1;  // The value here doesn't matter
			
			vector<Point> n = neighbors(candidate, filtered.size());
			candidates.insert(candidates.end(), n.begin(), n.end());
		}
	}
	
	return pixels;
}
	

/********** Private Methods ***************************************************/

bool filterBlue(Vec3b colors) {
	int diffRed = colors[BLUE] - colors[RED];
	int diffGreen = colors[BLUE] - colors[GREEN];
	
	return diffRed > 30 && diffGreen > 30;
}

bool filterGreen(Vec3b colors) {
	int diffRed = colors[GREEN] - colors[RED];
	int diffBlue = colors[GREEN] - colors[BLUE];
	
	return diffRed > 30 && diffBlue > 30;
}

void BlobDetector::filter(const Mat &src, Mat &dst, Color color) const {
	bool (*fptr)(Vec3b) = NULL;
	
	switch (color) {
	case COLOR_BLUE:
		fptr = &filterBlue;
		break;
	
	case COLOR_GREEN:
		fptr = &filterGreen;
		break;
	
	default:
		ERROR("Unknown color");
		return;
	}
	
	dst = Mat(src.size(), CV_8UC3);
	
	// TODO This should be more efficient
	for (int row = 0; row < src.size().height; row++) {
		for (int col = 0; col < src.size().width; col++) {
			for (int i = 0; i < 3; i++) {
				dst.at<Vec3b>(row, col)[i] = 0;
			}
		}
	}
	
	for (int row = 0; row < src.size().height; row++) {
		for (int col = 0; col < src.size().width; col++) {
			if ((*fptr)(src.at<Vec3b>(row, col))) {
				dst.at<Vec3b>(row, col)[BLUE] = 255;
			}
		}
	}
	
	cvtColor(dst, dst, CV_BGR2HSV);
}

vector<BlobDetector::Blob> BlobDetector::findBlobs(const Mat& rgbImage,
	Color color)
{
	std::vector<Blob> blobs;

	Mat hsvImage;
	filter(rgbImage, hsvImage, color);
	
	vector<KeyPoint> sbdResult;
	m_sbd.detect(hsvImage, sbdResult);
	
	// Iterates over all of the elements in sbdResult
	for (vector<KeyPoint>::const_iterator it = sbdResult.begin();
		it != sbdResult.end(); ++it)
	{
		Blob blob;
		blob.center = it->pt;
		blob.size = it->size;
		blob.color = color;
		blob.category = UNCATEGORIZED;
		
		blobs.push_back(blob);
	}
	
	return blobs;
}

/********** Local Static Functions ********************************************/

bool blobCompare(const BlobDetector::Blob &one, const BlobDetector::Blob &two)
{
	return one.size >= two.size;
}

// Take the euclidean of the difference between p1 and p2
double euclidean(Point p1, Point p2)
{ 
	return sqrt(pow((double) p2.x - p1.x, 2) + pow((double) p2.y - p1.y, 2));
}
